// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "kernels/funcs/npu_funcs.h"
#include "kernels/funcs/npu_op_runner.h"
#include "kernels/funcs/slice_utils.h"

namespace custom_kernel {

inline std::vector<int> get_new_shape_npu(
    const phi::CustomContext& dev_ctx,
    const std::vector<const phi::DenseTensor*>& list_new_shape_tensor) {
  // get tensor from
  std::vector<int> vec_new_shape;
  for (size_t i = 0; i < list_new_shape_tensor.size(); ++i) {
    auto tensor = list_new_shape_tensor[i];
    PADDLE_ENFORCE_EQ(tensor->dims() == phi::make_ddim({1}) ||
                          tensor->dims() == phi::make_ddim({}),
                      true,
                      phi::errors::InvalidArgument(
                          "The shape of dimension tensor should be [1] or [],"
                          "but received d%.",
                          tensor->dims()));
    if (tensor->dtype() == phi::DataType::INT64) {
      std::vector<int64_t> temp_vec(1);
      dev_ctx.Wait();
      TensorToVector(dev_ctx, *tensor, dev_ctx, &temp_vec);
      vec_new_shape.push_back(temp_vec[0]);
    } else if (tensor->dtype() == phi::DataType::INT32) {
      std::vector<int32_t> temp_vec(1);
      dev_ctx.Wait();
      TensorToVector(dev_ctx, *tensor, dev_ctx, &temp_vec);
      vec_new_shape.push_back(temp_vec[0]);
    }
  }

  return vec_new_shape;
}

template <typename T, typename Context>
void TransposeKernel(const Context& dev_ctx,
                     const phi::DenseTensor& x,
                     const std::vector<int>& axis,
                     phi::DenseTensor* out);

template <typename T, typename Context>
struct InterpolateFunction {
 public:
  explicit InterpolateFunction(const Context& dev_ctx) : dev_ctx(dev_ctx) {
    place = dev_ctx.GetPlace();
    stream = dev_ctx.stream();
  }
  void Arange(int n, phi::DenseTensor* x) {
    if (x->dtype() == phi::DataType::FLOAT16) {
      phi::DenseTensor x_fp32;
      phi::DenseTensorMeta x_fp32_meta = {phi::DataType::FLOAT32, x->dims()};
      x_fp32.set_meta(x_fp32_meta);
      dev_ctx.template Alloc<float>(&x_fp32);

      NpuOpRunner runner;
      // Fix refer to
      // https://gitee.com/ascend/modelzoo/issues/I6K3HN?from=project-issue
      runner.SetType("Range")
          .AddInput(dev_ctx,
                    std::move(std::vector<float>{static_cast<float>(0)}),
                    false)
          .AddInput(dev_ctx,
                    std::move(std::vector<float>{static_cast<float>(n)}),
                    false)
          .AddInput(dev_ctx,
                    std::move(std::vector<float>{static_cast<float>(1)}),
                    false)
          .AddOutput(x_fp32);
      runner.Run(stream);
      Cast(&x_fp32, x);
    } else {
      NpuOpRunner runner;
      // Fix refer to
      // https://gitee.com/ascend/modelzoo/issues/I6K3HN?from=project-issue
      runner.SetType("Range")
          .AddInput(dev_ctx,
                    std::move(std::vector<float>{static_cast<float>(0)}),
                    false)
          .AddInput(dev_ctx,
                    std::move(std::vector<float>{static_cast<float>(n)}),
                    false)
          .AddInput(dev_ctx,
                    std::move(std::vector<float>{static_cast<float>(1)}),
                    false)
          .AddOutput(*x);
      runner.Run(stream);
    }
  }
  void ReduceSum(const phi::DenseTensor* x,
                 phi::DenseTensor* y,
                 const std::vector<int>& dim,
                 bool keep_dims = true) {
    NpuOpRunner runner;
    runner.SetType("ReduceSum")
        .AddInput(*x)
        .AddInput(dev_ctx, std::move(dim))
        .AddOutput(*y)
        .AddAttrs({{"keep_dims", keep_dims}})
        .Run(dev_ctx.stream());
  }
  void Add(const phi::DenseTensor* x,
           const phi::DenseTensor* y,
           phi::DenseTensor* z) {
    const auto& runner = NpuOpRunner("AddV2", {*x, *y}, {*z}, {});
    runner.Run(stream);
  }
  void Adds(const phi::DenseTensor* x, float scalar, phi::DenseTensor* y) {
    const auto& runner = NpuOpRunner("Adds", {*x}, {*y}, {{"value", scalar}});
    runner.Run(stream);
  }
  void Mul(const phi::DenseTensor* x,
           const phi::DenseTensor* y,
           phi::DenseTensor* z) {
    const auto& runner = NpuOpRunner("Mul", {*x, *y}, {*z}, {});
    runner.Run(stream);
  }
  void Sub(const phi::DenseTensor* x,
           const phi::DenseTensor* y,
           phi::DenseTensor* z) {
    const auto& runner = NpuOpRunner("Sub", {*x, *y}, {*z}, {});
    runner.Run(stream);
  }
  void Cast(const phi::DenseTensor* x, phi::DenseTensor* y) {
    auto dst_dtype = ConvertToNpuDtype(y->dtype());
    const auto& runner = NpuOpRunner(
        "Cast", {*x}, {*y}, {{"dst_type", static_cast<int>(dst_dtype)}});
    runner.Run(stream);
  }
  void Gather(const phi::DenseTensor* x,
              const phi::DenseTensor* indices,
              const int axis,
              phi::DenseTensor* y) {
    EXEC_NPU_CMD(aclnnGatherV2, dev_ctx, *x, axis, *indices, *y);
  }
  void GatherWithTrans(const phi::DenseTensor* x,
                       const phi::DenseTensor* indices,
                       const int axis,
                       phi::DenseTensor* y) {
    // x: n, c, h, w
    PD_CHECK(axis == 3, "axis must be 3 for nchw input");
    auto x_dims = x->dims();
    auto y_dims = y->dims();
    phi::DenseTensor trans_x;
    trans_x.Resize({x_dims[0], x_dims[1], x_dims[3], x_dims[2]});
    std::vector<int> perm = {0, 1, 3, 2};

    // 1. transpose nchw -> ncwh
    custom_kernel::TransposeKernel<T, Context>(dev_ctx, *x, perm, &trans_x);

    // 2. gather ncwh
    phi::DenseTensor trans_y;
    trans_y.Resize({y_dims[0], y_dims[1], y_dims[3], y_dims[2]});
    dev_ctx.template Alloc<T>(&trans_y);
    int axis_gather = axis - 1;
    Gather(&trans_x, indices, axis_gather, &trans_y);

    // 3. tranpose ncwh -> nchw
    custom_kernel::TransposeKernel<T, Context>(dev_ctx, trans_y, perm, y);
  }
  void GatherGrad(const phi::DenseTensor* gy,
                  const phi::DenseTensor* indices,
                  const int axis,
                  phi::DenseTensor* gx) {
    //  1  gy swapaxis: axis & 0
    int len = (gy->dims()).size();
    std::vector<int> axis_swap(len);
    for (int i = 0; i < len; i++) {
      axis_swap[i] = i;
    }
    axis_swap[0] = axis;
    axis_swap[axis] = 0;
    auto y_new_shape = gy->dims();
    auto yt = y_new_shape[axis];
    y_new_shape[axis] = y_new_shape[0];
    y_new_shape[0] = yt;
    phi::DenseTensor gy_t;
    gy_t.Resize(y_new_shape);
    dev_ctx.template Alloc<T>(&gy_t);
    Transpose(dev_ctx, gy, &gy_t, axis_swap);

    //  2  scatter
    auto x_new_shape = gx->dims();
    auto xt = x_new_shape[axis];
    x_new_shape[axis] = x_new_shape[0];
    x_new_shape[0] = xt;
    phi::DenseTensor gx_zero;
    gx_zero.Resize(x_new_shape);
    dev_ctx.template Alloc<T>(&gx_zero);
    FillNpuTensorWithConstant<T>(&gx_zero, dev_ctx, static_cast<T>(0));
    gx_zero.Resize(x_new_shape);
    Scatter(dev_ctx, &gx_zero, indices, &gy_t, &gx_zero);

    //  3  gx swapaxis: axis, 0
    Transpose(dev_ctx, &gx_zero, gx, axis_swap);
  }
  void Scatter(const Context& dev_ctx,
               const phi::DenseTensor* x,
               const phi::DenseTensor* index,
               const phi::DenseTensor* updates,
               phi::DenseTensor* y) {
    const auto& runner =
        NpuOpRunner("ScatterNdAdd", {*x, *index, *updates}, {*y}, {});
    runner.Run(stream);
  }
  void Transpose(const Context& dev_ctx,
                 const phi::DenseTensor* x,
                 phi::DenseTensor* y,
                 const std::vector<int>& axis) {
    NpuOpRunner runner;
    runner.SetType("Transpose")
        .AddInput(*x)
        .AddInput(dev_ctx, std::move(axis))
        .AddOutput(*y)
        .Run(dev_ctx.stream());
  }
  void Muls(const phi::DenseTensor* x, float scalar, phi::DenseTensor* y) {
    const auto& runner = NpuOpRunner("Muls", {*x}, {*y}, {{"value", scalar}});
    runner.Run(stream);
  }
  void Maximum(const phi::DenseTensor* x,
               const phi::DenseTensor* y,
               phi::DenseTensor* z) {
    const auto& runner = NpuOpRunner("Maximum", {*x, *y}, {*z}, {});
    runner.Run(stream);
  }
  void Minimum(const phi::DenseTensor* x,
               const phi::DenseTensor* y,
               phi::DenseTensor* z) {
    const auto& runner = NpuOpRunner("Minimum", {*x, *y}, {*z}, {});
    runner.Run(stream);
  }
  void Floor(const phi::DenseTensor* x, phi::DenseTensor* y) {
    const auto& runner = NpuOpRunner("Floor", {*x}, {*y}, {});
    runner.Run(stream);
  }

 private:
  phi::Place place;
  aclrtStream stream;
  const Context& dev_ctx;
};

void InterpolateParamCompute(const float scale_h,
                             const float scale_w,
                             const bool align_corners,
                             const int align_mode,
                             const phi::DataLayout& data_layout,
                             const phi::DDim& indim,
                             const phi::DDim& outdim,
                             int* axis_h,
                             int* axis_w,
                             int* in_h,
                             int* in_w,
                             int* out_h,
                             int* out_w,
                             float* ratio_h,
                             float* ratio_w) {
  if (data_layout == phi::DataLayout::kNCHW) {
    *axis_h = 2;
    *axis_w = 3;
  } else {
    *axis_h = 1;
    *axis_w = 2;
  }
  *out_h = outdim[*axis_h];
  *out_w = outdim[*axis_w];
  *in_h = indim[*axis_h];
  *in_w = indim[*axis_w];
  *ratio_h = 0.0f;
  *ratio_w = 0.0f;
  if (*out_h > 1) {
    *ratio_h =
        align_corners
            ? static_cast<float>(*in_h - 1) / (*out_h - 1)
            : (scale_h > 0 ? 1 / scale_h : static_cast<float>(*in_h) / *out_h);
  }
  if (*out_w > 1) {
    *ratio_w =
        align_corners
            ? static_cast<float>(*in_w - 1) / (*out_w - 1)
            : (scale_w > 0 ? 1 / scale_w : static_cast<float>(*in_w) / *out_w);
  }
}

template <typename T, typename Context>
void BilinearParamTensorCompute(const Context& dev_ctx,
                                const phi::DataLayout& data_layout,
                                int in_h,
                                int in_w,
                                int out_h,
                                int out_w,
                                bool align_cond,
                                float ratio_h,
                                float ratio_w,
                                phi::DenseTensor* h0,
                                phi::DenseTensor* h1,
                                phi::DenseTensor* w0,
                                phi::DenseTensor* w1,
                                phi::DenseTensor* coef_h0,
                                phi::DenseTensor* coef_h1,
                                phi::DenseTensor* coef_w0,
                                phi::DenseTensor* coef_w1) {
  InterpolateFunction<T, Context> F(dev_ctx);
  auto place = dev_ctx.GetPlace();
  phi::DenseTensor _h0, _w0;
  _h0.Resize({out_h});
  _w0.Resize({out_w});
  dev_ctx.template Alloc<T>(&_h0);
  dev_ctx.template Alloc<T>(&_w0);
  F.Arange(out_h, &_h0);
  F.Arange(out_w, &_w0);
  if (align_cond) {
    F.Adds(&_h0, static_cast<float>(0.5), &_h0);
    F.Adds(&_w0, static_cast<float>(0.5), &_w0);
    F.Muls(&_h0, ratio_h, &_h0);
    F.Muls(&_w0, ratio_w, &_w0);
    F.Adds(&_h0, static_cast<float>(-0.5), &_h0);
    F.Adds(&_w0, static_cast<float>(-0.5), &_w0);
  } else {
    F.Muls(&_h0, ratio_h, &_h0);
    F.Muls(&_w0, ratio_w, &_w0);
  }

  phi::DenseTensor zero_t;
  phi::DenseTensor one_t;
  zero_t.Resize({1});
  one_t.Resize({1});
  dev_ctx.template Alloc<T>(&zero_t);
  dev_ctx.template Alloc<T>(&one_t);
  FillNpuTensorWithConstant<T>(&zero_t, dev_ctx, static_cast<T>(0));
  FillNpuTensorWithConstant<T>(&one_t, dev_ctx, static_cast<T>(1));
  F.Maximum(&_h0, &zero_t, &_h0);
  F.Maximum(&_w0, &zero_t, &_w0);

  phi::DenseTensor _h0_floor, _w0_floor;
  _h0_floor.Resize({out_h});
  _w0_floor.Resize({out_w});
  dev_ctx.template Alloc<T>(&_h0_floor);
  dev_ctx.template Alloc<T>(&_w0_floor);
  F.Floor(&_h0, &_h0_floor);
  F.Floor(&_w0, &_w0_floor);
  F.Cast(&_h0_floor, h0);
  F.Cast(&_w0_floor, w0);

  phi::DenseTensor one_int;
  one_int.Resize({1});
  dev_ctx.template Alloc<T>(&one_int);
  FillNpuTensorWithConstant<int>(&one_int, dev_ctx, static_cast<int>(1));
  F.Add(h0, &one_int, h1);
  F.Add(w0, &one_int, w1);
  phi::DenseTensor t_max_h, t_max_w;
  t_max_h.Resize({1});
  t_max_w.Resize({1});
  dev_ctx.template Alloc<T>(&t_max_h);
  dev_ctx.template Alloc<T>(&t_max_w);
  FillNpuTensorWithConstant<int>(&t_max_h, dev_ctx, static_cast<int>(in_h - 1));
  FillNpuTensorWithConstant<int>(&t_max_w, dev_ctx, static_cast<int>(in_w - 1));
  F.Minimum(h1, &t_max_h, h1);
  F.Minimum(w1, &t_max_w, w1);

  F.Sub(&_h0, &_h0_floor, coef_h1);
  F.Sub(&_w0, &_w0_floor, coef_w1);
  F.Sub(&one_t, coef_h1, coef_h0);
  F.Sub(&one_t, coef_w1, coef_w0);

  if (data_layout == phi::DataLayout::kNCHW) {
    coef_h0->Resize({out_h, 1});
    coef_h1->Resize({out_h, 1});
  } else {
    coef_h0->Resize({out_h, 1, 1});
    coef_h1->Resize({out_h, 1, 1});
    coef_w0->Resize({out_w, 1});
    coef_w1->Resize({out_w, 1});
  }
}

template <typename T, typename Context>
void BilinearFwdAclnn(const Context& dev_ctx,
                      const phi::DenseTensor* input,
                      phi::DenseTensor* output,
                      const float scale_h,
                      const float scale_w,
                      const bool align_corners,
                      const int align_mode,
                      const phi::DataLayout& data_layout) {
  auto outdim = output->dims();
  double h = 1;
  double w = 1;
  std::vector<int64_t> outsize;
  for (int i = 2; i < outdim.size(); i++) {
    outsize.push_back(outdim[i]);
  }
  EXEC_NPU_CMD(aclnnUpsampleBilinear2d,
               dev_ctx,
               *input,
               outsize,
               align_corners,
               h,
               w,
               *output);
}

template <typename T, typename Context>
void BilinearFwdNpu(const Context& dev_ctx,
                    const phi::DenseTensor* input,
                    phi::DenseTensor* output,
                    const float scale_h,
                    const float scale_w,
                    const bool align_corners,
                    const int align_mode,
                    const phi::DataLayout& data_layout) {
  auto outdim = output->dims();
  auto indim = input->dims();

  if (data_layout == phi::DataLayout::NCHW) {
    int indim_h = indim[2];
    int indim_w = indim[3];
    int outdim_h = outdim[2];
    int outdim_w = outdim[3];

    bool aclnn_flag = false;
    if (scale_h == -1 && scale_w == -1) {
      aclnn_flag = true;
    } else if (static_cast<float>(outdim_h) / static_cast<float>(indim_h) ==
                   scale_h &&
               static_cast<float>(outdim_w) / static_cast<float>(indim_w) ==
                   scale_w) {
      aclnn_flag = true;
    }

    if (align_corners == true) {
      return custom_kernel::BilinearFwdAclnn<T, Context>(dev_ctx,
                                                         input,
                                                         output,
                                                         scale_h,
                                                         scale_w,
                                                         align_corners,
                                                         align_mode,
                                                         data_layout);
    } else if (aclnn_flag && align_mode == 0) {
      if (outdim_h != 1 && outdim_w != 1) {
        return custom_kernel::BilinearFwdAclnn<T, Context>(dev_ctx,
                                                           input,
                                                           output,
                                                           scale_h,
                                                           scale_w,
                                                           align_corners,
                                                           align_mode,
                                                           data_layout);
      } else if (outdim_h == 1 && outdim_w == 1) {
        return custom_kernel::BilinearFwdAclnn<T, Context>(dev_ctx,
                                                           input,
                                                           output,
                                                           scale_h,
                                                           scale_w,
                                                           true,
                                                           align_mode,
                                                           data_layout);
      }
    }
  }

  InterpolateFunction<T, Context> F(dev_ctx);
  auto place = dev_ctx.GetPlace();

  int axis_h, axis_w;
  int out_h, out_w, in_h, in_w;
  float ratio_h, ratio_w;
  InterpolateParamCompute(scale_h,
                          scale_w,
                          align_corners,
                          align_mode,
                          data_layout,
                          indim,
                          outdim,
                          &axis_h,
                          &axis_w,
                          &in_h,
                          &in_w,
                          &out_h,
                          &out_w,
                          &ratio_h,
                          &ratio_w);

  phi::DenseTensor h0, h1, w0, w1;
  h0.Resize({out_h});
  h1.Resize({out_h});
  w0.Resize({out_w});
  w1.Resize({out_w});
  dev_ctx.template Alloc<int>(&h0);
  dev_ctx.template Alloc<int>(&h1);
  dev_ctx.template Alloc<int>(&w0);
  dev_ctx.template Alloc<int>(&w1);
  phi::DenseTensor coef_h0, coef_h1, coef_w0, coef_w1;
  coef_h0.Resize({out_h});
  coef_h1.Resize({out_h});
  coef_w0.Resize({out_w});
  coef_w1.Resize({out_w});
  dev_ctx.template Alloc<T>(&coef_h0);
  dev_ctx.template Alloc<T>(&coef_h1);
  dev_ctx.template Alloc<T>(&coef_w0);
  dev_ctx.template Alloc<T>(&coef_w1);
  bool align_cond = align_mode == 0 && !align_corners;
  BilinearParamTensorCompute<T>(dev_ctx,
                                data_layout,
                                in_h,
                                in_w,
                                out_h,
                                out_w,
                                align_cond,
                                ratio_h,
                                ratio_w,
                                &h0,
                                &h1,
                                &w0,
                                &w1,
                                &coef_h0,
                                &coef_h1,
                                &coef_w0,
                                &coef_w1);

  phi::DenseTensor input_gather_h0, input_gather_h1;
  auto dim_gather_h = indim;
  dim_gather_h[axis_h] = out_h;
  input_gather_h0.Resize(dim_gather_h);
  input_gather_h1.Resize(dim_gather_h);
  dev_ctx.template Alloc<T>(&input_gather_h0);
  dev_ctx.template Alloc<T>(&input_gather_h1);
  F.Gather(input, &h0, axis_h, &input_gather_h0);
  F.Gather(input, &h1, axis_h, &input_gather_h1);

  F.Mul(&input_gather_h0, &coef_h0, &input_gather_h0);
  F.Mul(&input_gather_h1, &coef_h1, &input_gather_h1);
  phi::DenseTensor out_x4;
  out_x4.Resize({4, outdim[0], outdim[1], outdim[2], outdim[3]});
  dev_ctx.template Alloc<T>(&out_x4);
  phi::DenseTensor input_gather_h0_w0 = custom_kernel::Slice(out_x4, 0, 1);
  input_gather_h0_w0.Resize({outdim[0], outdim[1], outdim[2], outdim[3]});
  phi::DenseTensor input_gather_h0_w1 = custom_kernel::Slice(out_x4, 1, 2);
  input_gather_h0_w1.Resize({outdim[0], outdim[1], outdim[2], outdim[3]});
  phi::DenseTensor input_gather_h1_w0 = custom_kernel::Slice(out_x4, 2, 3);
  input_gather_h1_w0.Resize({outdim[0], outdim[1], outdim[2], outdim[3]});
  phi::DenseTensor input_gather_h1_w1 = custom_kernel::Slice(out_x4, 3, 4);
  input_gather_h1_w1.Resize({outdim[0], outdim[1], outdim[2], outdim[3]});
  if (axis_w == 3) {  // nchw, gather along the last axis;
    F.GatherWithTrans(&input_gather_h0, &w0, axis_w, &input_gather_h0_w0);
    F.GatherWithTrans(&input_gather_h0, &w1, axis_w, &input_gather_h0_w1);
    F.GatherWithTrans(&input_gather_h1, &w0, axis_w, &input_gather_h1_w0);
    F.GatherWithTrans(&input_gather_h1, &w1, axis_w, &input_gather_h1_w1);
  } else {
    F.Gather(&input_gather_h0, &w0, axis_w, &input_gather_h0_w0);
    F.Gather(&input_gather_h0, &w1, axis_w, &input_gather_h0_w1);
    F.Gather(&input_gather_h1, &w0, axis_w, &input_gather_h1_w0);
    F.Gather(&input_gather_h1, &w1, axis_w, &input_gather_h1_w1);
  }
  F.Mul(&input_gather_h0_w0, &coef_w0, &input_gather_h0_w0);
  F.Mul(&input_gather_h0_w1, &coef_w1, &input_gather_h0_w1);
  F.Mul(&input_gather_h1_w0, &coef_w0, &input_gather_h1_w0);
  F.Mul(&input_gather_h1_w1, &coef_w1, &input_gather_h1_w1);
  F.ReduceSum(&out_x4, output, std::vector<int>{0}, false);
}

template <typename T, typename Context>
void BilinearBwdAclnn(const Context& dev_ctx,
                      const phi::DenseTensor* gout,
                      phi::DenseTensor* gin,
                      const float scale_h,
                      const float scale_w,
                      const bool align_corners,
                      const int align_mode,
                      const phi::DataLayout& data_layout) {
  auto indim = gin->dims();
  auto outdim = gout->dims();
  double h = 1;
  double w = 1;

  std::vector<int64_t> outputsize;
  for (int i = 2; i < outdim.size(); i++) {
    outputsize.push_back(outdim[i]);
  }
  std::vector<int64_t> inputsize;
  for (int i = 0; i < indim.size(); i++) {
    inputsize.push_back(indim[i]);
  }
  // dev_ctx.template Alloc<T>(gin);
  EXEC_NPU_CMD(aclnnUpsampleBilinear2dBackward,
               dev_ctx,
               *gout,
               outputsize,
               inputsize,
               align_corners,
               h,
               w,
               *gin);
}

template <typename T, typename Context>
void BilinearBwdNpu(const Context& dev_ctx,
                    const phi::DenseTensor* gout,
                    phi::DenseTensor* gin,
                    const float scale_h,
                    const float scale_w,
                    const bool align_corners,
                    const int align_mode,
                    const phi::DataLayout& data_layout) {
  auto outdim = gout->dims();
  auto indim = gin->dims();

  if (data_layout == phi::DataLayout::NCHW) {
    int indim_h = indim[2];
    int indim_w = indim[3];
    int outdim_h = outdim[2];
    int outdim_w = outdim[3];

    bool aclnn_flag = false;
    if (scale_h == -1 && scale_w == -1) {
      aclnn_flag = true;
    } else if (static_cast<float>(outdim_h) / static_cast<float>(indim_h) ==
                   scale_h &&
               static_cast<float>(outdim_w) / static_cast<float>(indim_w) ==
                   scale_w) {
      aclnn_flag = true;
    }

    if (align_corners == true) {
      return custom_kernel::BilinearBwdAclnn<T, Context>(dev_ctx,
                                                         gout,
                                                         gin,
                                                         scale_h,
                                                         scale_w,
                                                         align_corners,
                                                         align_mode,
                                                         data_layout);
    } else if (aclnn_flag && align_mode == 0) {
      if (outdim_h != 1 && outdim_w != 1) {
        return custom_kernel::BilinearBwdAclnn<T, Context>(dev_ctx,
                                                           gout,
                                                           gin,
                                                           scale_h,
                                                           scale_w,
                                                           align_corners,
                                                           align_mode,
                                                           data_layout);
      } else if (outdim_h == 1 && outdim_w == 1) {
        return custom_kernel::BilinearBwdAclnn<T, Context>(dev_ctx,
                                                           gout,
                                                           gin,
                                                           scale_h,
                                                           scale_w,
                                                           true,
                                                           align_mode,
                                                           data_layout);
      }
    }
  }

  InterpolateFunction<T, Context> F(dev_ctx);
  auto place = dev_ctx.GetPlace();

  int axis_h, axis_w;
  int out_h, out_w, in_h, in_w;
  float ratio_h, ratio_w;
  InterpolateParamCompute(scale_h,
                          scale_w,
                          align_corners,
                          align_mode,
                          data_layout,
                          indim,
                          outdim,
                          &axis_h,
                          &axis_w,
                          &in_h,
                          &in_w,
                          &out_h,
                          &out_w,
                          &ratio_h,
                          &ratio_w);

  phi::DenseTensor h0, h1, w0, w1;
  h0.Resize({out_h});
  h1.Resize({out_h});
  w0.Resize({out_w});
  w1.Resize({out_w});
  dev_ctx.template Alloc<int>(&h0);
  dev_ctx.template Alloc<int>(&h1);
  dev_ctx.template Alloc<int>(&w0);
  dev_ctx.template Alloc<int>(&w1);
  phi::DenseTensor coef_h0, coef_h1, coef_w0, coef_w1;
  coef_h0.Resize({out_h});
  coef_h1.Resize({out_h});
  coef_w0.Resize({out_w});
  coef_w1.Resize({out_w});
  dev_ctx.template Alloc<T>(&coef_h0);
  dev_ctx.template Alloc<T>(&coef_h1);
  dev_ctx.template Alloc<T>(&coef_w0);
  dev_ctx.template Alloc<T>(&coef_w1);
  bool align_cond = align_mode == 0 && !align_corners;
  BilinearParamTensorCompute<T, Context>(dev_ctx,
                                         data_layout,
                                         in_h,
                                         in_w,
                                         out_h,
                                         out_w,
                                         align_cond,
                                         ratio_h,
                                         ratio_w,
                                         &h0,
                                         &h1,
                                         &w0,
                                         &w1,
                                         &coef_h0,
                                         &coef_h1,
                                         &coef_w0,
                                         &coef_w1);
  phi::DenseTensor gy_w0, gy_w1;
  gy_w0.Resize(outdim);
  gy_w1.Resize(outdim);
  dev_ctx.template Alloc<T>(&gy_w0);
  dev_ctx.template Alloc<T>(&gy_w1);
  F.Mul(gout, &coef_w0, &gy_w0);
  F.Mul(gout, &coef_w1, &gy_w1);

  auto dim_gather_h = indim;
  dim_gather_h[axis_h] = out_h;
  phi::DenseTensor g_gather_w0, g_gather_w1;
  g_gather_w0.Resize(dim_gather_h);
  g_gather_w1.Resize(dim_gather_h);
  dev_ctx.template Alloc<T>(&g_gather_w0);
  dev_ctx.template Alloc<T>(&g_gather_w1);

  w0.Resize({out_w, 1});
  w1.Resize({out_w, 1});
  F.GatherGrad(&gy_w0, &w0, axis_w, &g_gather_w0);
  F.GatherGrad(&gy_w1, &w1, axis_w, &g_gather_w1);

  F.Add(&g_gather_w0, &g_gather_w1, &g_gather_w0);
  F.Mul(&g_gather_w0, &coef_h1, &g_gather_w1);
  F.Mul(&g_gather_w0, &coef_h0, &g_gather_w0);

  phi::DenseTensor gx_0, gx_1;
  gx_0.Resize(indim);
  gx_1.Resize(indim);
  dev_ctx.template Alloc<T>(&gx_0);
  dev_ctx.template Alloc<T>(&gx_1);
  h0.Resize({out_h, 1});
  h1.Resize({out_h, 1});
  F.GatherGrad(&g_gather_w0, &h0, axis_h, &gx_0);
  F.GatherGrad(&g_gather_w1, &h1, axis_h, &gx_1);

  F.Add(&gx_0, &gx_1, gin);
}

template <typename T, typename Context>
void AclopNearestFwdNpu(const Context& dev_ctx,
                        const phi::DenseTensor& input,
                        const std::vector<int64_t>& out_shape,
                        const bool align_corners,
                        phi::DenseTensor* output) {
  auto stream = dev_ctx.stream();

  NpuOpRunner runner;
  runner.SetType("ResizeNearestNeighborV2")
      .AddInput(input)
      .AddInput(dev_ctx, std::vector<int64_t>(out_shape))
      .AddOutput(*output)
      .AddAttr("align_corners", align_corners)
      .AddAttr("half_pixel_centers", false);
  runner.Run(stream);
}

template <typename T, typename Context>
void NearestFwdNpu(const Context& dev_ctx,
                   const phi::DenseTensor& input,
                   const std::vector<int64_t>& out_shape,
                   const bool align_corners,
                   phi::DenseTensor* output) {
  DO_COMPATIBILITY(aclnnUpsampleNearest2d,
                   (custom_kernel::AclopNearestFwdNpu<T, Context>(
                       dev_ctx, input, out_shape, align_corners, output)));

  EXEC_NPU_CMD(aclnnUpsampleNearest2d, dev_ctx, input, out_shape, *output);
}

template <typename T, typename Context>
void InterpolateKernel(
    const Context& dev_ctx,
    const phi::DenseTensor& x,
    const paddle::optional<phi::DenseTensor>& out_size,
    const paddle::optional<std::vector<const phi::DenseTensor*>>& size_tensor,
    const paddle::optional<phi::DenseTensor>& scale_tensor,
    const std::string& data_layout_str,
    int out_d,
    int out_h,
    int out_w,
    const std::vector<float>& scale,
    const std::string& interp_method,
    bool align_corners,
    int align_mode,
    phi::DenseTensor* out) {
  auto input = x;
  auto input_dims = input.dims();
  PADDLE_ENFORCE_EQ(
      input_dims.size(),
      4UL,
      phi::errors::External("NPU Interpolate Kernel only support 4-D Tensor."));

  phi::DataLayout data_layout = StringToDataLayout(data_layout_str);
  int n, c, in_d, in_h, in_w;
  ExtractNCDWH(input_dims, data_layout, &n, &c, &in_d, &in_h, &in_w);

  float scale_h = -1;
  float scale_w = -1;

  // Priority: SizeTensor > OutSize > Scale > scale > out_h & out_w
  if (size_tensor && size_tensor->size() > 0) {
    auto output_get = get_new_shape_npu(dev_ctx, size_tensor.get());
    if (output_get.size() <= 2) {
      out_h = output_get[0];
      out_w = output_get[1];
    } else {
      out_h = output_get[0];
      out_w = output_get[1];
      out_d = output_get[2];
    }
  } else {
    if (scale_tensor) {
      auto scale_data =
          get_new_data_from_tensor<float>(dev_ctx, scale_tensor.get_ptr());
      if (scale_data.size() > 1) {
        scale_h = scale_data[0];
        scale_w = scale_data[1];
      } else {
        scale_h = scale_data[0];
        scale_w = scale_data[0];
      }
      PADDLE_ENFORCE_EQ(
          scale_w > 0,
          true,
          phi::errors::InvalidArgument(
              "The scale_w in input 'Scale' Tensor of Operator(interpolate) "
              "should be greater than 0, but received value is %d.",
              scale_w));
      PADDLE_ENFORCE_EQ(
          scale_h > 0,
          true,
          phi::errors::InvalidArgument(
              "The scale_h in input 'Scale' Tensor of Operator(interpolate) "
              "should be greater than 0, but received value is %d.",
              scale_h));
    } else {
      if (scale.size() > 1) {
        scale_h = scale[0];
        scale_w = scale[1];

        PADDLE_ENFORCE_EQ(
            scale_w > 0,
            true,
            phi::errors::InvalidArgument(
                "The scale_w in Attr(scale) of Operator(interpolate) "
                "should be greater than 0, but received value is %d.",
                scale_w));
        PADDLE_ENFORCE_EQ(
            scale_h > 0,
            true,
            phi::errors::InvalidArgument(
                "The scale_h in Attr(scale) of Operator(interpolate) "
                "should be greater than 0, but received value is %d.",
                scale_h));
      }
    }
    if (scale_h > 0. && scale_w > 0.) {
      out_h = static_cast<int>(in_h * scale_h);
      out_w = static_cast<int>(in_w * scale_w);
    }
    if (out_size) {
      auto size_data =
          get_new_data_from_tensor<int>(dev_ctx, out_size.get_ptr());
      // phi::DenseTensor sizes;
      // TensorCopy(dev_ctx, *out_size, true, &sizes, phi::CPUPlace());
      // auto size_data = sizes.data<int>();
      out_h = size_data[0];
      out_w = size_data[1];
    }
  }
  PADDLE_ENFORCE_GT(out_h,
                    0,
                    phi::errors::InvalidArgument(
                        "out_h in Attr(out_shape) of Op(interpolate) "
                        "should be greater than 0."));
  PADDLE_ENFORCE_GT(out_w,
                    0,
                    phi::errors::InvalidArgument(
                        "out_w in Attr(out_shape) of Op(interpolate) "
                        "should be greater than 0."));
  phi::DDim dim_out;
  if (data_layout == phi::DataLayout::kNCHW) {
    dim_out = {n, c, out_h, out_w};
  } else {
    dim_out = {n, out_h, out_w, c};
  }

  phi::DenseTensorMeta out_meta = {out->dtype(), dim_out, data_layout};
  out->set_meta(out_meta);
  dev_ctx.template Alloc<T>(out);

  if (in_h == out_h && in_w == out_w) {
    TensorCopy(dev_ctx, input, false, out);
    return;
  }

  // To-do(qili93): need to support bilineare, try ResizeD
  // Add bilineare by zhulei
  if ("nearest" == interp_method) {
    std::vector<int64_t> out_shape = {out_h, out_w};
    NearestFwdNpu<T, Context>(dev_ctx, input, out_shape, align_corners, out);
  } else if ("bilinear" == interp_method) {
    BilinearFwdNpu<T, Context>(dev_ctx,
                               &input,
                               out,
                               scale_h,
                               scale_w,
                               align_corners,
                               align_mode,
                               data_layout);
  }
}

template <typename T, typename Context>
void AclopNearestBwdNpu(const Context& dev_ctx,
                        const phi::DenseTensor& output_grad,
                        const std::vector<int64_t>& in_shape,
                        const bool align_corners,
                        phi::DenseTensor* input_grad) {
  auto stream = dev_ctx.stream();

  NpuOpRunner runner;
  runner.SetType("ResizeNearestNeighborV2Grad")
      .AddInput(output_grad)
      .AddInput(dev_ctx, std::vector<int64_t>(in_shape))
      .AddOutput(*input_grad)
      .AddAttr("align_corners", align_corners)
      .AddAttr("half_pixel_centers", false);
  runner.Run(stream);
}

template <typename T, typename Context>
void NearestBwdNpu(const Context& dev_ctx,
                   const phi::DenseTensor& output_grad,
                   const std::vector<int64_t>& out_shape,
                   const std::vector<int64_t>& in_shape,
                   const double scales_h,
                   const double scales_w,
                   const bool align_corners,
                   phi::DenseTensor* input_grad) {
  DO_COMPATIBILITY(
      aclnnUpsampleNearest2dBackward,
      (custom_kernel::AclopNearestBwdNpu<T, Context>(
          dev_ctx, output_grad, in_shape, align_corners, input_grad)));

  std::vector<int64_t> input_dims = phi::vectorize<int64_t>(input_grad->dims());
  EXEC_NPU_CMD(aclnnUpsampleNearest2dBackward,
               dev_ctx,
               output_grad,
               out_shape,
               input_dims,
               scales_h,
               scales_w,
               *input_grad);
}

template <typename T, typename Context>
void InterpolateGradKernel(
    const Context& dev_ctx,
    const phi::DenseTensor& x,
    const paddle::optional<phi::DenseTensor>& out_size,
    const paddle::optional<std::vector<const phi::DenseTensor*>>& size_tensor,
    const paddle::optional<phi::DenseTensor>& scale_tensor,
    const phi::DenseTensor& out_grad,
    const std::string& data_layout_str,
    int out_d,
    int out_h,
    int out_w,
    const std::vector<float>& scale,
    const std::string& interp_method,
    bool align_corners,
    int align_mode,
    phi::DenseTensor* dx) {
  auto input = x;
  auto input_grad = dx;
  auto output_grad = out_grad;

  auto stream = dev_ctx.stream();

  phi::DataLayout data_layout = StringToDataLayout(data_layout_str);
  int n, c, in_d, in_h, in_w;
  ExtractNCDWH(input.dims(), data_layout, &n, &c, &in_d, &in_h, &in_w);

  float scale_h = -1;
  float scale_w = -1;

  // Priority: SizeTensor > OutSize > Scale > scale > out_h & out_w
  if (size_tensor && size_tensor->size() > 0) {
    auto output_get = get_new_shape_npu(dev_ctx, size_tensor.get());
    out_h = output_get[0];
    out_w = output_get[1];
  } else if (out_size) {
    auto out_size_data =
        get_new_data_from_tensor<int>(dev_ctx, out_size.get_ptr());
    out_h = out_size_data[0];
    out_w = out_size_data[1];

    scale_h = static_cast<float>(out_h) / static_cast<float>(in_h);
    scale_w = static_cast<float>(out_w) / static_cast<float>(in_w);
  } else {
    if (scale_tensor) {
      auto scale_data =
          get_new_data_from_tensor<float>(dev_ctx, scale_tensor.get_ptr());
      if (scale_data.size() > 1) {
        scale_h = scale_data[0];
        scale_w = scale_data[1];
      } else {
        scale_w = scale_data[0];
        scale_h = scale_data[0];
      }
      PADDLE_ENFORCE_EQ(
          scale_w > 0,
          true,
          phi::errors::InvalidArgument(
              "The scale_w in input 'Scale' Tensor of Operator(interpolate) "
              "should be greater than 0, but received value is %d.",
              scale_w));
      PADDLE_ENFORCE_EQ(
          scale_h > 0,
          true,
          phi::errors::InvalidArgument(
              "The scale_h in input 'Scale' Tensor of Operator(interpolate) "
              "should be greater than 0, but received value is %d.",
              scale_h));
    } else {
      if (scale.size() > 1) {
        scale_h = scale[0];
        scale_w = scale[1];
        PADDLE_ENFORCE_EQ(
            scale_w > 0,
            true,
            phi::errors::InvalidArgument(
                "The scale_w in Attr(scale) of Operator(interpolate) "
                "should be greater than 0, but received value is %d.",
                scale_w));
        PADDLE_ENFORCE_EQ(
            scale_h > 0,
            true,
            phi::errors::InvalidArgument(
                "The scale_h in Attr(scale) of Operator(interpolate) "
                "should be greater than 0, but received value is %d.",
                scale_h));
      }
    }
    if (scale_h > 0. && scale_w > 0.) {
      out_h = static_cast<int>(in_h * scale_h);
      out_w = static_cast<int>(in_w * scale_w);
    }
  }

  phi::DDim dim_grad;
  if (data_layout == phi::DataLayout::kNCHW) {
    dim_grad = {n, c, in_h, in_w};
  } else {
    dim_grad = {n, in_h, in_w, c};
  }

  phi::DenseTensorMeta input_grad_meta = {input.dtype(), dim_grad, data_layout};
  input_grad->set_meta(input_grad_meta);
  dev_ctx.template Alloc<T>(input_grad);

  if (in_h == out_h && in_w == out_w) {
    TensorCopy(dev_ctx, output_grad, false, input_grad);
    return;
  }

  // To-do(qili93): need to support bilineare, try ResizeGradD
  if ("nearest" == interp_method) {
    std::vector<int64_t> out_shape = {out_h, out_w};
    std::vector<int64_t> in_shape = {in_h, in_w};
    NearestBwdNpu<T, Context>(dev_ctx,
                              output_grad,
                              out_shape,
                              in_shape,
                              scale_h,
                              scale_w,
                              align_corners,
                              input_grad);
  } else if ("bilinear" == interp_method) {
    BilinearBwdNpu<T, Context>(dev_ctx,
                               &output_grad,
                               input_grad,
                               scale_h,
                               scale_w,
                               align_corners,
                               align_mode,
                               data_layout);
  }
}

template <typename T, typename Context>
void BilinearInterpKernel(
    const Context& ctx,
    const phi::DenseTensor& x,
    const paddle::optional<phi::DenseTensor>& out_size,
    const paddle::optional<std::vector<const phi::DenseTensor*>>& size_tensor,
    const paddle::optional<phi::DenseTensor>& scale_tensor,
    const std::string& data_layout,
    int out_d,
    int out_h,
    int out_w,
    const std::vector<float>& scale,
    const std::string& interp_method,
    bool align_corners,
    int align_mode,
    phi::DenseTensor* output) {
  InterpolateKernel<T, Context>(ctx,
                                x,
                                out_size,
                                size_tensor,
                                scale_tensor,
                                data_layout,
                                out_d,
                                out_h,
                                out_w,
                                scale,
                                interp_method,
                                align_corners,
                                align_mode,
                                output);
}

template <typename T, typename Context>
void NearestInterpKernel(
    const Context& ctx,
    const phi::DenseTensor& x,
    const paddle::optional<phi::DenseTensor>& out_size,
    const paddle::optional<std::vector<const phi::DenseTensor*>>& size_tensor,
    const paddle::optional<phi::DenseTensor>& scale_tensor,
    const std::string& data_layout,
    int out_d,
    int out_h,
    int out_w,
    const std::vector<float>& scale,
    const std::string& interp_method,
    bool align_corners,
    int align_mode,
    phi::DenseTensor* output) {
  InterpolateKernel<T, Context>(ctx,
                                x,
                                out_size,
                                size_tensor,
                                scale_tensor,
                                data_layout,
                                out_d,
                                out_h,
                                out_w,
                                scale,
                                interp_method,
                                align_corners,
                                align_mode,
                                output);
}

template <typename T, typename Context>
void BilinearInterpGradKernel(
    const Context& dev_ctx,
    const phi::DenseTensor& x,
    const paddle::optional<phi::DenseTensor>& out_size,
    const paddle::optional<std::vector<const phi::DenseTensor*>>& size_tensor,
    const paddle::optional<phi::DenseTensor>& scale_tensor,
    const phi::DenseTensor& out_grad,
    const std::string& data_layout,
    int out_d,
    int out_h,
    int out_w,
    const std::vector<float>& scale,
    const std::string& interp_method,
    bool align_corners,
    int align_mode,
    phi::DenseTensor* x_grad) {
  InterpolateGradKernel<T, Context>(dev_ctx,
                                    x,
                                    out_size,
                                    size_tensor,
                                    scale_tensor,
                                    out_grad,
                                    data_layout,
                                    out_d,
                                    out_h,
                                    out_w,
                                    scale,
                                    interp_method,
                                    align_corners,
                                    align_mode,
                                    x_grad);
}

template <typename T, typename Context>
void NearestInterpGradKernel(
    const Context& dev_ctx,
    const phi::DenseTensor& x,
    const paddle::optional<phi::DenseTensor>& out_size,
    const paddle::optional<std::vector<const phi::DenseTensor*>>& size_tensor,
    const paddle::optional<phi::DenseTensor>& scale_tensor,
    const phi::DenseTensor& out_grad,
    const std::string& data_layout,
    int out_d,
    int out_h,
    int out_w,
    const std::vector<float>& scale,
    const std::string& interp_method,
    bool align_corners,
    int align_mode,
    phi::DenseTensor* x_grad) {
  InterpolateGradKernel<T, Context>(dev_ctx,
                                    x,
                                    out_size,
                                    size_tensor,
                                    scale_tensor,
                                    out_grad,
                                    data_layout,
                                    out_d,
                                    out_h,
                                    out_w,
                                    scale,
                                    interp_method,
                                    align_corners,
                                    align_mode,
                                    x_grad);
}

}  // namespace custom_kernel

PD_REGISTER_PLUGIN_KERNEL(nearest_interp,
                          npu,
                          ALL_LAYOUT,
                          custom_kernel::NearestInterpKernel,
                          float,
                          phi::dtype::float16) {
  kernel->InputAt(1).SetBackend(phi::Backend::ALL_BACKEND);
  kernel->InputAt(2).SetBackend(phi::Backend::ALL_BACKEND);
  kernel->InputAt(3).SetBackend(phi::Backend::ALL_BACKEND);
}

PD_REGISTER_PLUGIN_KERNEL(nearest_interp_grad,
                          npu,
                          ALL_LAYOUT,
                          custom_kernel::NearestInterpGradKernel,
                          float,
                          phi::dtype::float16) {
  kernel->InputAt(2).SetBackend(phi::Backend::ALL_BACKEND);
  kernel->InputAt(3).SetBackend(phi::Backend::ALL_BACKEND);
}

PD_REGISTER_PLUGIN_KERNEL(bilinear_interp,
                          npu,
                          ALL_LAYOUT,
                          custom_kernel::BilinearInterpKernel,
                          float,
                          phi::dtype::float16) {
  kernel->InputAt(1).SetBackend(phi::Backend::ALL_BACKEND);
  kernel->InputAt(2).SetBackend(phi::Backend::ALL_BACKEND);
  kernel->InputAt(3).SetBackend(phi::Backend::ALL_BACKEND);
}

PD_REGISTER_PLUGIN_KERNEL(bilinear_interp_grad,
                          npu,
                          ALL_LAYOUT,
                          custom_kernel::BilinearInterpGradKernel,
                          float,
                          phi::dtype::float16) {
  kernel->InputAt(2).SetBackend(phi::Backend::ALL_BACKEND);
  kernel->InputAt(3).SetBackend(phi::Backend::ALL_BACKEND);
}
